# This script does the main out-of-sample (projection) counterfactual
# calculations from Section 6.3 of the paper

# See "do_counterfactuals_insample.py" for the insample calculation
# See "do_counterfactuals_boot.py" for the bootstrap calculation
# Use "plot_counterfactuals.py" to make Figures 7 and 8

# Jacob Moscona and Karthik Sastry
# This version: last edited on September 28, 2022


##
## Importing libraries and files
##

import numpy as np
import pandas as pd
import os

wd = '/home/karthik/Dropbox (MIT)/climate_crops/QJE_Submission/Replication/Data/'
os.chdir(wd)

# File with the shocks
historical_file = 'county_shocks.csv'

# File with two extra pieces: total area in 2017 and price index
extra_file = 'county_extra_for_counterfactual.csv'

# Which scenarios to consider
# Default behavior: loop over all of them
# path_list =  ['45','60','85']
# tfinal = [2050,2090]

path_list =  ['45','60','85']
tfinal = [2050,2090]



##
## Options
##

## GDD, panel, 1950 to 2010
fe_file = 'gdd_panel_1950_2010_fe.xls'
beta = -0.000232
gamma = -0.000823
phi = 9.12e-08
k = 7.828


##
## Main Calculation
##


for path in path_list:
    for tf in tfinal:

        future_file = 'county_shocks' + '_' + path + '_' + str(tf) + '_2012a.csv'
        # dta_future = pd.read_csv(future_file)
        dta_future = pd.read_csv('/home/karthik/Dropbox (MIT)/climate_crops/Data/' + 'county_level_May2020_projection' + str(tf) + '_' + path + '_area2012.csv')
    
    
        t0 = 2010 # Initial decade
        t1 = tf # Final decade
        
        metric = 'gdd' # Name of the metric being used
        
        
        # Historical values of own, leave one out, etc.
        dta = pd.read_csv(historical_file)
        dta = dta.drop(columns = [metric + x + '2010' for x in ['_own_','_loo_']])
        
        dta = dta.merge(dta_future, how = 'left', on = ['STATE_FIPS','CNTY_FIPS'])
        
        # Values of fixed effect 
        dta_fe = pd.read_excel('../Data/' + fe_file)
        dta_fe.columns = ['year','CNTY_FIPS','STATE_FIPS','area','area_init','ife','tfe']
        
        # Merge IFE
        dta = dta.merge(dta_fe.loc[dta_fe.year == t0,['CNTY_FIPS','STATE_FIPS','ife']], on =
                        ['CNTY_FIPS','STATE_FIPS'])
        # Merge TFE0
        dta = dta.merge(dta_fe.loc[dta_fe.year == t0,['CNTY_FIPS','STATE_FIPS','tfe']], on =
                        ['CNTY_FIPS','STATE_FIPS'])
        dta = dta.rename(columns = {'tfe':'tfe0'})
    
        # Merge TFE from history
        dta = dta.merge(dta_fe.loc[dta_fe.year == (t0-50),['CNTY_FIPS','STATE_FIPS','tfe']], on =
                        ['CNTY_FIPS','STATE_FIPS'])    
        dta = dta.rename(columns = {'tfe':'tfeold'})
        
        scale = (tf-2010)/50
        dta['tfe1'] = dta['tfe0'] + scale * (dta['tfe0']-dta['tfeold']) # Constant growth in time FE
        
        # Merge initial area
        dta = dta.merge(dta_fe.loc[dta_fe.year == t0,['CNTY_FIPS','STATE_FIPS','area_init']], on =
                        ['CNTY_FIPS','STATE_FIPS'])    
        
        # Merge extra columns
        dta_extra = pd.read_csv(extra_file)
        dta = dta.merge(dta_extra, on =['CNTY_FIPS','STATE_FIPS'])
    
        
        rho_i = dta['ife'].values
        alpha = [[],[]]
        alpha[0] = dta['tfe0'].values+k
        alpha[1] = dta['tfe1'].values+k
        area = dta.total_area_acres_2017.values

        ##
        ## Main function
        ##
            
        def fitted_values(own,loo,interact,alpha,rho):
            return rho + alpha + beta * own + gamma * loo + phi * (interact)
            
        def calculate_land_values(own0,loo0,own1,loo1,rho_i,offset=0):
        
            fit0 = fitted_values(own0,loo0,own0*loo0,alpha[0],rho_i)
            fit1 = fitted_values(own1,loo1,own1*loo1,alpha[1],rho_i)
            
            int_noi = own1 * (loo0 + offset) # counterfactual eco interact
            fit_noi = fitted_values(own1,loo1,int_noi,alpha[1],rho_i) # no innovation
            
            fit_nocc = fitted_values(own0,loo0,own0*loo0,alpha[1],rho_i)
            
            return(fit0,fit1,fit_noi,fit_nocc)
            
        ##
        ## Calculation
        ##
            
        
        # Trim extreme outliers, as described in the paper
        dta[metric + '_own_' + str(t1)] = np.minimum(dta[metric + '_own_' + str(t1)],15000)
        dta[metric + '_own_' + str(t0)] = np.minimum(dta[metric + '_own_' + str(t0)],15000)
        
        # Treating the zero as zero; offset = 0
        fit0,fit1,fit_noi,fit_nocc = calculate_land_values(
                    dta[metric + '_own_' + str(t0)],
                    dta[metric + '_loo_' + str(t0)],
                    dta[metric + '_own_' + str(t1)],
                    dta[metric + '_loo_' + str(t1)],
                    rho_i,
                    offset = 0
                )
        
        ##
        ## Total summary
        ##
        
        def logmean(x):
            return np.log(np.nanmean(np.exp(x))) 
            
        
        percent_saved = 100 * (logmean(fit_noi) - logmean(fit_nocc)) # Units = percentage difference
        percent_damage = 100 * (logmean(fit1) - logmean(fit_nocc)) # Units = percentage difference
        awt = area
        
        total_val = np.nansum(awt* np.exp(fit1))
        noi_val = np.nansum(awt* np.exp(fit_noi))
        nocc_val = np.nansum(awt* np.exp(fit_nocc))

        x = awt* np.exp(fit_nocc) - awt* np.exp(fit1)    
        loss_with_innovation = np.nansum(x)
        
        y  = awt* np.exp(fit_nocc) - awt* np.exp(fit_noi)    
        loss_without_innovation = np.nansum(y)
        
        
        
        innovation_saved = loss_without_innovation - loss_with_innovation
        big_kahuna = 100 * innovation_saved / loss_without_innovation
        percent_damage = 100 * loss_with_innovation / np.nansum(awt* np.exp(fit_nocc))
        percent_damage_noi = 100 * loss_without_innovation / np.nansum(awt* np.exp(fit_nocc))
        
        print('The RCP is ' + path + '\n')
        print('The damage is ' + str(round(percent_damage,2)) + ' percent\n')
        print('The damage without innovation is ' + str(round(percent_damage_noi,2)) + ' percent\n')
        print('The percent saved is ' + str(round(big_kahuna,2)) + ' percent\n')
    
        dollar_savings = innovation_saved/1e9 / (1.03 ** (t1-2010)) # 3 pct inflation = current value
        print('The dollar savings are '+ str(round(dollar_savings,2)) + ' billion\n')
        print('The dollar value is '+ str(round(total_val/1e12,4)) + ' trillion\n')
        print('The dollar value, no I, is '+ str(round(noi_val/1e12,4)) + ' trillion\n')
        print('The dollar value, no CC, is '+ str(round(nocc_val/1e12,4)) + ' trillion\n')


